"""
频谱分析仪模块
基于RsInstrument_FSW_iq_transfer_display.py重构，负责FSW频谱分析仪的控制和IQ数据采集
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from typing import Dict, Any, Optional, Tuple
from RsInstrument import *
from logger import Logger


class SpectrumAnalyzer:
    """FSW频谱分析仪控制器"""
    
    def __init__(self, config: Dict[str, Any], logger: Optional[Logger] = None):
        """
        初始化频谱分析仪
        
        Args:
            config: FSW仪器配置
            logger: 日志记录器
        """
        self.config = config
        self.logger = logger or Logger()
        self.fsw = None
        self.iq_data = None
        # 记录关键信息以便绘制频谱
        self.center_freq: Optional[float] = None
        self.sample_rate: Optional[float] = None
        self.capture_time: Optional[float] = None
        
    def connect(self) -> bool:
        """
        连接FSW频谱分析仪
        
        Returns:
            连接是否成功
        """
        try:
            ip = self.config.get('ip', '192.168.31.240')
            self.logger.info(f"正在连接FSW频谱分析仪: {ip}")
            
            # 检查RsInstrument模块版本
            RsInstrument.assert_minimum_version('1.53.0')
            
            # 建立连接
            self.fsw = RsInstrument(f'TCPIP::{ip}::hislip0', True, False)
            
            # 获取仪器信息
            idn_string = self.fsw.idn_string
            driver_version = self.fsw.driver_version
            self.logger.info(f"FSW仪器ID: {idn_string}")
            self.logger.info(f"驱动版本: {driver_version}")
            
            # 重置仪器到已知状态
            self.fsw.clear_status()
            self.fsw.reset()
            self.logger.info("FSW仪器已重置")
            
            return True
            
        except Exception as e:
            self.logger.error(f"连接FSW频谱分析仪失败: {e}")
            return False
    
    def configure_basic_settings(self, center_freq: float, ref_level: float) -> bool:
        """
        配置基本设置
        
        Args:
            center_freq: 中心频率 (Hz)
            ref_level: 参考电平 (dBm)
            
        Returns:
            配置是否成功
        """
        try:
            self.logger.info(f"配置基本设置: 中心频率={center_freq/1e9:.3f} GHz, 参考电平={ref_level} dBm")
            
            # 设置参考电平
            self.fsw.write_str(f'DISPlay:WINDow:TRACe:Y:SCALe:RLEVel {ref_level}')
            
            # 设置中心频率
            self.fsw.write_str(f'FREQ:CENT {center_freq}')
            # 保存用于绘图
            self.center_freq = center_freq
            
            self.logger.info("基本设置配置完成")
            return True
            
        except Exception as e:
            self.logger.error(f"配置基本设置失败: {e}")
            return False
    
    def configure_iq_analyzer(self, sample_rate: float, capture_time: float) -> bool:
        """
        配置IQ分析仪模式
        
        Args:
            sample_rate: 采样率 (Hz)
            capture_time: 捕获时间 (秒)
            
        Returns:
            配置是否成功
        """
        try:
            self.logger.info(f"配置IQ分析仪模式: 采样率={sample_rate/1e6:.2f} MS/s, 捕获时间={capture_time} s")
            
            # 创建IQ分析仪模式
            self.fsw.write_with_opc('INST:CRE:NEW IQ, "IQ Analyzer"')
            
            # 设置为单次扫描
            self.fsw.write_str('INIT:CONT OFF')
            
            # 设置采样率
            self.fsw.write_str(f'TRAC:IQ:SRAT {sample_rate}')

            # 设置捕获时间
            self.fsw.write_str(f'SENS:SWE:TIME {capture_time}')
            # 保存用于绘图
            self.sample_rate = sample_rate
            self.capture_time = capture_time
            
            self.logger.info("IQ分析仪模式配置完成")
            return True
            
        except Exception as e:
            self.logger.error(f"配置IQ分析仪模式失败: {e}")
            return False
    
    def configure_trigger(self, trigger_config: Dict[str, Any], capture_time: float) -> bool:
        """
        配置触发器
        
        Args:
            trigger_config: 触发器配置
            capture_time: 捕获时间
            
        Returns:
            配置是否成功
        """
        try:
            trigger_source = trigger_config.get('source', 'EXT')
            trigger_slope = trigger_config.get('slope', 'POS')
            trigger_level = trigger_config.get('level', 1.0)
            trigger_timeout = trigger_config.get('timeout', 10)
            
            self.logger.info(f"配置触发器: 源={trigger_source}, 斜率={trigger_slope}, 电平={trigger_level}")
            
            # 设置VISA超时
            visa_timeout = int(capture_time * 1000) + 20000 + int(trigger_timeout * 1000)
            self.fsw.visa_timeout = visa_timeout
            self.fsw.opc_timeout = visa_timeout
            
            # 设置触发源
            self.fsw.write_str(f'TRIG:SOUR {trigger_source}')
            
            # 根据触发源设置相应的电平和斜率
            if trigger_source.upper() == 'EXT':
                self.fsw.write_str(f'TRIG:SLOP {trigger_slope}')
                self.fsw.write_str(f'TRIG:LEVel {trigger_level}V')
                self.logger.info(f"使用外部触发器 (EXT), 电平: {trigger_level}V")
                
            elif trigger_source.upper() == 'RFPO':
                self.fsw.write_str(f'TRIG:LEVel:RFPower {trigger_level}DBM')
                self.logger.info(f"使用RF功率触发器 (RFPO), 电平: {trigger_level}dBm")
            
            self.logger.info("触发器配置完成")
            return True
            
        except Exception as e:
            self.logger.error(f"配置触发器失败: {e}")
            return False
    
    def wait_for_trigger_and_acquisition(self, trigger_config: Dict[str, Any]) -> bool:
        """
        等待触发事件并完成数据采集
        
        Args:
            trigger_config: 触发器配置
            
        Returns:
            采集是否成功
        """
        try:
            trigger_source = trigger_config.get('source', 'EXT')
            trigger_level = trigger_config.get('level', 1.0)
            trigger_slope = trigger_config.get('slope', 'POS')
            
            self.logger.info("等待触发事件...")
            self.logger.info(f"触发器源: {trigger_source}")
            
            if trigger_source.upper() == 'RFPO':
                self.logger.info(f"触发电平: {trigger_level}dBm")
            elif trigger_source.upper() == 'EXT':
                self.logger.info(f"触发电平: {trigger_level}V")
            
            self.logger.info(f"触发斜率: {trigger_slope}")
            
            # 启动单次扫描，等待触发
            self.fsw.write_with_opc('INIT:IMM')
            
            # 等待操作完成
            self.logger.info('等待触发事件和数据采集...')
            self.fsw.query_opc()
            
            self.logger.info('✓ 触发检测到，数据采集完成')
            return True
            
        except Exception as e:
            self.logger.warning(f'⚠ 警告: 触发等待超时或发生错误: {e}')
            self.logger.warning('尝试使用可用数据继续...')
            
            # 检查系统状态
            try:
                sys_status = self.fsw.query_str(':SYST:ERR?')
                self.logger.info(f'系统状态: {sys_status}')
            except:
                pass
            
            return False
    
    def prepare_for_trigger(self, trigger_config: Dict[str, Any]) -> bool:
        """
        准备触发器状态，进入就绪状态但不立即启动采集
        适用于先准备频谱仪，再开启信号源的工作流程
        
        Args:
            trigger_config: 触发器配置
            
        Returns:
            准备是否成功
        """
        try:
            trigger_source = trigger_config.get('source', 'EXT')
            trigger_level = trigger_config.get('level', 1.0)
            trigger_slope = trigger_config.get('slope', 'POS')
            
            self.logger.info("准备触发器状态...")
            self.logger.info(f"触发器源: {trigger_source}")
            
            if trigger_source.upper() == 'RFPO':
                self.logger.info(f"触发电平: {trigger_level}dBm")
            elif trigger_source.upper() == 'EXT':
                self.logger.info(f"触发电平: {trigger_level}V")
            
            self.logger.info(f"触发斜率: {trigger_slope}")
            
            # 配置触发器但不立即启动扫描
            # 仅设置触发器参数，等待后续的INIT:IMM命令
            self.logger.info("频谱仪已准备就绪，等待触发事件...")
            
            return True
            
        except Exception as e:
            self.logger.error(f"准备触发器状态失败: {e}")
            return False
    
    def start_triggered_acquisition(self) -> bool:
        """
        启动已准备的触发采集
        在prepare_for_trigger()之后调用
        
        Returns:
            采集是否成功
        """
        try:
            # 启动单次扫描，等待触发
            self.fsw.write_with_opc('INIT:IMM')
            
            # 等待操作完成
            self.logger.info('等待触发事件和数据采集...')
            self.fsw.query_opc()
            
            self.logger.info('✓ 触发检测到，数据采集完成')
            return True
            
        except Exception as e:
            self.logger.warning(f'⚠ 警告: 触发等待超时或发生错误: {e}')
            self.logger.warning('尝试使用可用数据继续...')
            
            # 检查系统状态
            try:
                sys_status = self.fsw.query_str(':SYST:ERR?')
                self.logger.info(f'系统状态: {sys_status}')
            except:
                pass
            
            return False
    
    def acquire_iq_data(self) -> bool:
        """
        采集IQ数据
        
        Returns:
            采集是否成功
        """
        try:
            self.logger.info("正在读取IQ数据...")
            
            # 设置数据格式为32位实数
            self.fsw.write_str('FORM REAL,32')
            
            # 设置IQ数据格式
            self.fsw.write_with_opc('TRAC:IQ:DATA:FORM IQP')
            
            # 读取二进制数据
            raw_data = np.array(self.fsw.query_bin_or_ascii_float_list('TRAC:IQ:DATA:MEM?'))
            
            self.logger.info(f"✓ 成功读取 {len(raw_data)} 个数据点")
            
            # 转换为复数IQ数据
            self.iq_data = raw_data[0::2] + 1j * raw_data[1::2]
            self.logger.info(f"✓ 创建了 {len(self.iq_data)} 个复数IQ样本")

            # 直接绘制频谱图
            try:
                self._plot_spectrum(self.iq_data)
            except Exception as e:
                self.logger.warning(f"绘制频谱图失败: {e}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"✗ 读取IQ数据时出错: {e}")
            self.iq_data = None
            return False
    
    def get_iq_data(self) -> Optional[np.ndarray]:
        """
        获取IQ数据
        
        Returns:
            IQ数据数组，如果未采集则返回None
        """
        return self.iq_data

    # -------------------------
    # 绘图与分析辅助
    # -------------------------
    def _plot_spectrum(self, iq: np.ndarray) -> None:
        """
        绘制单次两边频谱（dBm/Hz），参考 RsInstrument_FSW_iq_transfer_display.py 的实现。
        需在调用前已通过 configure_basic_settings / configure_iq_analyzer 设置中心频率与采样率。
        """
        if iq is None or len(iq) == 0:
            self.logger.warning("无IQ数据，跳过频谱绘制")
            return
        if self.sample_rate is None or self.center_freq is None:
            self.logger.warning("缺少采样率或中心频率，跳过频谱绘制")
            return

        fs = float(self.sample_rate)
        cf = float(self.center_freq)

        N = len(iq)
        window_type = 'hann'
        # 使用下一个2的幂作为FFT点数，保证效率与频率分辨率
        nfft = 1 << (N - 1).bit_length()
        self.logger.info(f"绘制频谱: N={N}, fs={fs/1e6:.3f} MS/s, nfft={nfft}")

        # 周期图估计功率谱密度（V^2/Hz），然后换算为 dBm/Hz（50 Ohm）
        f_psd, Pxx = signal.periodogram(
            iq,
            fs=fs,
            window=window_type,
            nfft=nfft,
            scaling='density',
            detrend=False,
            return_onesided=False
        )

        f_psd = np.fft.fftshift(f_psd)
        Pxx = np.fft.fftshift(Pxx)
        f_offset = f_psd
        f_abs = f_offset + cf
        Pxx_dbm = 10 * np.log10(Pxx / 50.0 * 1000.0 + 1e-12)

        # 画图
        plt.figure(figsize=(12, 5))
        plt.plot(f_offset / 1e6, Pxx_dbm, lw=1.1)
        plt.xlabel('Frequency Offset (MHz)')
        plt.ylabel('Power Spectral Density (dBm/Hz)')
        plt.title(f'Triggered Signal Spectrum\nCenter: {cf/1e9:.3f} GHz')
        plt.grid(True, alpha=0.25)
        plt.tight_layout()
        plt.show()

        # 日志输出一些统计值
        peak_idx = int(np.argmax(Pxx_dbm))
        self.logger.info(
            f"峰值: {Pxx_dbm[peak_idx]:.2f} dBm/Hz @ 偏移 {f_offset[peak_idx]/1e6:.3f} MHz (绝对 {f_abs[peak_idx]/1e9:.6f} GHz)"
        )
    
    def get_iq_raw_data(self) -> Optional[np.ndarray]:
        """
        获取原始IQ数据（交错格式）
        
        Returns:
            原始IQ数据数组，如果未采集则返回None
        """
        if self.iq_data is None:
            return None
        
        # 重构交错格式数据
        raw_data = np.zeros(2 * len(self.iq_data))
        raw_data[0::2] = self.iq_data.real
        raw_data[1::2] = self.iq_data.imag
        
        return raw_data
    
    def acquire_iq_with_config(self, center_freq: float, ref_level: float, 
                              sample_rate: float, capture_time: float,
                              trigger_config: Dict[str, Any]) -> bool:
        """
        使用指定配置完整采集IQ数据
        
        Args:
            center_freq: 中心频率 (Hz)
            ref_level: 参考电平 (dBm)
            sample_rate: 采样率 (Hz)
            capture_time: 捕获时间 (秒)
            trigger_config: 触发器配置
            
        Returns:
            采集是否成功
        """
        try:
            self.logger.info("=" * 60)
            self.logger.info("开始IQ数据采集")
            self.logger.info("=" * 60)
            
            # 配置基本设置
            if not self.configure_basic_settings(center_freq, ref_level):
                return False
            
            # 配置IQ分析仪
            if not self.configure_iq_analyzer(sample_rate, capture_time):
                return False
            
            # 配置触发器
            if not self.configure_trigger(trigger_config, capture_time):
                return False
            
            # 等待触发和采集
            if not self.wait_for_trigger_and_acquisition(trigger_config):
                return False
            
            # 读取IQ数据
            if not self.acquire_iq_data():
                return False
            
            self.logger.info("=" * 60)
            self.logger.info("IQ数据采集完成")
            self.logger.info("=" * 60)
            
            return True
            
        except Exception as e:
            self.logger.error(f"IQ数据采集失败: {e}")
            return False
    
    def get_instrument_info(self) -> Dict[str, Any]:
        """
        获取仪器信息
        
        Returns:
            仪器信息字典
        """
        if self.fsw is None:
            return {}
        
        try:
            return {
                'idn': self.fsw.idn_string,
                'driver_version': self.fsw.driver_version,
                'connected': True
            }
        except Exception:
            return {'connected': False}
    
    def disconnect(self) -> None:
        """断开与FSW的连接"""
        try:
            if self.fsw:
                self.fsw.close()
                self.fsw = None
                self.logger.info("FSW频谱分析仪连接已关闭")
                
        except Exception as e:
            self.logger.error(f"断开FSW连接时出错: {e}")
    
    def __del__(self):
        """析构函数，确保连接被正确关闭"""
        self.disconnect()
